library(data.table)
library(mefa)
library(doParallel)
library(emayili)
library(magrittr)
library(ggplot2)
library(Hmisc)
library(stats)
library(lfe)

source('helper_functions.r')


doFCEstimation <- function(estimation = 'estimation_final') {
  
  
  estimation.data <- loadEstimation(estimation)
  
  # Get market data
  market.data <- merge(estimation.data$market.data[,c('MARKET_ID','j','CBSA','YEAR','LENDER','N_LENDERS','MARKET_SIZE','SHARE_MODEL'),with=F],
                         estimation.data$mr[,c('j','MARKET_ID','LOAN_SIZE','MARKUP_MODEL')],by=c('MARKET_ID','j'))[LENDER == 'n']
  
  # Work out per-lender variable profits
  market.data[,VARIABLE_PROFIT := MARKET_SIZE * SHARE_MODEL/N_LENDERS * LOAN_SIZE * (MARKUP_MODEL / 100)/1000 ] # in k dollars, per lender
  market.data <- market.data[VARIABLE_PROFIT > 0]
  
  
  N = 815
  
  log_likelihood <- function(x) {
    print(x)
    m = x[1]
    s = x[2]
    #N = x[3]
    
    temp.data <- market.data[,c('N_LENDERS','VARIABLE_PROFIT','MARKET_SIZE'),with=F]
    temp.data[,market_log_likelihood :=   max(0,(N - N_LENDERS)) * log(1 - pnorm( log(VARIABLE_PROFIT),mean = m,sd = s)) + N_LENDERS * log(pnorm( log(VARIABLE_PROFIT),mean = m,sd = s)) ]
    
    val = -1 * weighted.mean(temp.data$market_log_likelihood,w=temp.data$MARKET_SIZE)
    # All lenders with costs above profit will not enter; all lenders with profit above costs will enter
    print(val)
  }
  
  
  toSolve <- optim(par = c(6.2,5.3),fn = log_likelihood,lower = c(1,1),upper = c(13,13))
  m.0 = toSolve$par[1]
  s.0 = toSolve$par[2]
  
  # Bootstrap it
  nB = 100
  all.ids <- unique(market.data$MARKET_ID)
  registerDoParallel(cores = 10)
  bootstraps <- foreach(ii = 1:nB, .combine = rbind) %dopar% {
    
    sampled.ids <- sample(x = all.ids,size = length(all.ids)/2,replace = T)
    current.data <- market.data[MARKET_ID %in% sampled.ids]
    
    log_likelihood <- function(x) {
      m = x[1]
      s = x[2]
      
      temp.data <- current.data[,c('N_LENDERS','VARIABLE_PROFIT','MARKET_SIZE'),with=F]
      temp.data[,market_log_likelihood :=   max(0,(N - N_LENDERS)) * log(1 - pnorm( log(VARIABLE_PROFIT),mean = m,sd = s)) + N_LENDERS * log(pnorm( log(VARIABLE_PROFIT),mean = m,sd = s)) ]
      
      val = -1 * weighted.mean(temp.data$market_log_likelihood,w=temp.data$MARKET_SIZE)
      # All lenders with costs above profit will not enter; all lenders with profit above costs will enter
      return(val)
    }
    
    toSolve <- optim(par = c(m.0,s.0),fn = log_likelihood,lower = c(1,1),upper = c(13,13))
    m = toSolve$par[1]
    s = toSolve$par[2]
    toRet = data.table(m = m,s = s)
  }
  
  result = data.table(variables = c('mu','sigma'), values = c(m.0,s.0), se = c(sd(bootstraps$m),sd(bootstraps$s)))
  
  
  # To get number of lenders---
  # For an individual lender, given variable profits, the probability that that lender is in is P(log(F) < log(VP)) = phi(log(VP); mu, sigma)
  market.data[,E.lenders := N * pnorm(log(VARIABLE_PROFIT),m.0,s.0)] # -- essentially the mode...
  aa <- lm(N_LENDERS ~ E.lenders,data=market.data,w = market.data$MARKET_SIZE)
  summary(aa)
  mean(market.data$E.lenders)
  mean(market.data$N_LENDERS)
}


doSupplyEstimation <- function(estimation = 'estimation_final') {
  
  # Load data
  estimation.data <- loadSupplyEstimationData(estimation)
  bank.data       <- fread('../data/final/capitalization_data_sep9.csv')
  
  
  # GMM Objective
  objective <- function(x,return.all = F) {
    # Initialize parameters
    pars <- list(psi = x[1],sig_eq = x[2],sig_gse = x[3])
    #pars <- list(psi = x[1],sig_eq = 1,sig_gse = 1)
    
    print(pars)
    
    # Calculate it
    mc <- calculateMC(estimation.data,bank.data,supply.pars = pars)
    
    # Objective function
    val <- sum(mc$moments$value^2)
    
    # Print some diagnostics
    print(val)
    print(mc$moments)
    
    if(!return.all) {
      return(val)
    } else {
      return(list(val = val,moments = mc$moments,mc = mc,pars = pars))
    }
  }
  
  ss <- optim(par = c(1.05,.25,.25),fn = objective,lower = c(1,.005,.005),upper = c(1.5,.75,.75),method = 'L-BFGS-B',control = list(lmm = 10,parscale=c(1,1,1)))
  result <- objective(ss$par,return.all = T)
  
  # Bootstrap estimation
  nB = 50
  all.ids <- unique(estimation.data$MARKET_ID)
  registerDoParallel(cores = 10)
  bootstraps <- foreach(ii = 1:nB, .combine = rbind) %dopar% {
    
    sampled.ids <- sample(x = all.ids,size = length(all.ids)/2,replace = T)
    current.data <- estimation.data[MARKET_ID %in% sampled.ids]
    
    objective <- function(x) {
      # Initialize parameters
      pars <- list(psi = x[1],sig_eq = x[2],sig_gse = x[3])
      
      # Calculate it
      mc <- calculateMC(current.data,bank.data,supply.pars = pars)
      
      # Objective function
      val <- sum(mc$moments$value^2,na.rm=T)
      return(val)
    }
    
    solution <- optim(par = c(result$pars$psi[1]+runif(1)/25-.5/25,result$pars$sig_eq[1]+runif(1)/25-.5/25,result$pars$sig_gse[1]+runif(1)/25-.5/25),fn = objective,lower = c(1,.005,.005),upper = c(1.5,.75,.75),method = 'L-BFGS-B',control = list(lmm = 10,parscale=c(1,1,1)))
    toRet <- data.table(sig_gse = solution$par[3],sig_re = solution$par[2],psi =solution$par[1])
  }

  
  # Get data to write.
  # Panel A: Financing costs. It's years + intercepts.
  table.A <- rbind(
    data.table(year = 2010,value = result$mc$linear$coefficients['(Intercept)'], se = summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']),
    data.table(year = 2011,value = result$mc$linear$coefficients['as.factor(YEAR)2011'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2011','Std. Error']^2)),
    data.table(year = 2012,value = result$mc$linear$coefficients['as.factor(YEAR)2012'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2012','Std. Error']^2)),
    data.table(year = 2013,value = result$mc$linear$coefficients['as.factor(YEAR)2013'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2013','Std. Error']^2)),
    data.table(year = 2014,value = result$mc$linear$coefficients['as.factor(YEAR)2014'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2014','Std. Error']^2)),
    data.table(year = 2015,value = result$mc$linear$coefficients['as.factor(YEAR)2015'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2015','Std. Error']^2)),
    data.table(year = 2016,value = result$mc$linear$coefficients['as.factor(YEAR)2016'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2016','Std. Error']^2)),
    data.table(year = 2017,value = result$mc$linear$coefficients['as.factor(YEAR)2017'] + result$mc$linear$coefficients['(Intercept)'], se = sqrt(summary(result$mc$linear)$coefficients['(Intercept)','Std. Error']^2 + summary(result$mc$linear)$coefficients['as.factor(YEAR)2017','Std. Error']^2))
  )
  
  table.B <- rbind(
    data.table(parameter = 'Bank, purchase (baseline)', value = 0, se = 0),
    data.table(parameter = 'Bank, refi',value = result$mc$linear$coefficients['PURPOSEr'],se = summary(result$mc$linear)$coefficients['PURPOSEr','Std. Error']),
    data.table(parameter = 'Shadow bank, non-fintech, purchase',value = result$mc$linear$coefficients['LENDERn'],se = summary(result$mc$linear)$coefficients['LENDERn','Std. Error']),
    data.table(parameter = 'Shadow bank, non-fintech, refi',value = result$mc$linear$coefficients['LENDERn'] + result$mc$linear$coefficients['PURPOSEr'] + result$mc$linear$coefficients['PURPOSEr:LENDERn'] ,se = sqrt(summary(result$mc$linear)$coefficients['LENDERn','Std. Error']^2 + summary(result$mc$linear)$coefficients['PURPOSEr','Std. Error']^2 + summary(result$mc$linear)$coefficients['PURPOSEr:LENDERn','Std. Error']^2)),
    data.table(parameter = 'Shadow bank, fintech, purchase',value = result$mc$linear$coefficients['LENDERf'],se = summary(result$mc$linear)$coefficients['LENDERf','Std. Error']),
    data.table(parameter = 'Shadow bank, fintech, refi',value = result$mc$linear$coefficients['LENDERf'] + result$mc$linear$coefficients['PURPOSEr'] + result$mc$linear$coefficients['PURPOSEr:LENDERf'] ,se = sqrt(summary(result$mc$linear)$coefficients['LENDERf','Std. Error']^2 + summary(result$mc$linear)$coefficients['PURPOSEr','Std. Error']^2 + summary(result$mc$linear)$coefficients['PURPOSEr:LENDERf','Std. Error']^2))
  )
  
  table.C <- rbind(
    data.table(parameter = 'High FICO, conforming (baseline)', value = 0, se = 0),
    data.table(parameter = 'High FICO, jumbo', value = result$mc$linear$coefficients['JUMBOj'],se = summary(result$mc$linear)$coefficients['JUMBOj','Std. Error']),
    data.table(parameter = 'Low FICO, conforming', value = result$mc$linear$coefficients['FICOl'], se = summary(result$mc$linear)$coefficients['FICOl','Std. Error']),
    data.table(parameter = 'Low FICO, jumbo', value = result$mc$linear$coefficients['FICOl'] + result$mc$linear$coefficients['JUMBOj'] + result$mc$linear$coefficients['FICOl:JUMBOj'], se = sqrt(summary(result$mc$linear)$coefficients['FICOl','Std. Error']^2 + summary(result$mc$linear)$coefficients['JUMBOj','Std. Error']^2 + summary(result$mc$linear)$coefficients['FICOl:JUMBOj','Std. Error']^2) )  
  )
  
  table.D <- rbind(
    data.table(parameter = 'sig_gse', value = result$pars[3], se = sd(bootstraps$sig_gse)),
    data.table(parameter = 'sig_re', value = result$pars[2], se = sd(bootstraps$sig_re)),
    data.table(parameter = 'psi', value = result$pars[1], se = sd(bootstraps$psi))
  )
  
  
  # Save the result.
  save(list = ls(),file = paste0('supply_estimation_results/',estimation,'.rsave'))
  
  
  
  plots <- plotCosts(result)
  plots[,Type := 'Bank conforming']
  plots[JUMBO == 'j',Type := 'Bank jumbo']
  plots[LENDER == 'n',Type := 'Shadow bank']
  
  ggplot(plots) + geom_line(aes(x=CAP_BANK,y=mc_nonlin + mc_lin,group = Type,linetype = Type)) + 
    geom_histogram(data=bank.data[BANK_CR < .2 & BANK_CR > .09],aes(x=BANK_CR,y=30*..count../sum(..count..)),alpha=.5,bins = 12) + theme_bw() + xlab('Bank capitalization') + ylab('Marginal cost') +
    scale_color_brewer(palette = 'Dark2') + theme(legend.position = c(.75,.75),legend.background = element_blank()) + coord_cartesian(xlim = c(.095,.195))
  ggsave(paste0('estimation_results/',estimation,'_costs.png'),height=4,width=6,units = 'in')
}


plotCosts <- function(estimation_result) {
  
  CAP = seq(0.09,.2,by = .0001)
  ECAP = CAP - .06
  
  params <- estimation_result$pars
  
  fake.data <- rbind(
    data.table(LENDER = 'b' , JUMBO = 'j', ECAP_BANK = ECAP,CAP_BANK = CAP,YEAR = 2017, PURPOSE = 'p',FICO = 'l',xi_j = .5,xi_g = .2),
    data.table(LENDER = 'b' , JUMBO = 'c', ECAP_BANK = ECAP,CAP_BANK = CAP,YEAR = 2017, PURPOSE = 'p',FICO = 'l',xi_j = .5,xi_g = .2),
    data.table(LENDER = 'n' , JUMBO = 'c', ECAP_BANK = ECAP,CAP_BANK = CAP,YEAR = 2017, PURPOSE = 'p',FICO = 'l',xi_j = .5,xi_g = .2))
  
  fake.data[,mc_jumbo_nonlin_balance      := CAP_BANK^2 * params$sig_eq * params$psi * xi_j * (ECAP_BANK)^(-1 * (params$psi + 1))]
  fake.data[,mc_conforming_nonlin_balance := CAP_BANK^2 * params$sig_eq * params$psi * xi_g * (ECAP_BANK)^(-1 * (params$psi + 1))]
  fake.data[,mc_conforming_nonlin_gse          := params$sig_gse]
  fake.data[LENDER == 'b' & JUMBO == 'j',mc_nonlin := mc_jumbo_nonlin_balance]
  fake.data[LENDER == 'b' & JUMBO == 'c',mc_nonlin := pmin(mc_conforming_nonlin_balance,mc_conforming_nonlin_gse)]
  fake.data[LENDER != 'b'               ,mc_nonlin := mc_conforming_nonlin_gse]
  
  
  fake.data[,mc_lin := predict(estimation_result$mc$linear,fake.data)]

  return(fake.data)
    
}
